package top.yinjinbiao.arcface.service.impl;

import cn.hutool.core.collection.CollectionUtil;
import com.arcsoft.face.*;
import com.arcsoft.face.enums.DetectMode;
import com.arcsoft.face.enums.DetectOrient;
import com.arcsoft.face.toolkit.ImageInfo;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import top.yinjinbiao.arcface.domain.entity.FacePerson;
import top.yinjinbiao.arcface.domain.dto.UserCompareInfo;
import top.yinjinbiao.arcface.factory.FaceEngineFactory;
import top.yinjinbiao.arcface.properties.ArcsoftProperties;
import top.yinjinbiao.arcface.common.exception.CustomException;
import top.yinjinbiao.arcface.common.exception.ErrorCodeEnum;
import top.yinjinbiao.arcface.service.FaceEngineService;

import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;

/**
 * @Description
 * @Author yin.jinbiao
 * @Date 2021/9/27 13:43
 * @Version 1.0
 */
@Service
@Slf4j
public class FaceEngineServiceImpl implements FaceEngineService {

    @Autowired
    private ArcsoftProperties arcsoftProperties;

    private ExecutorService compareExecutorService;

    /**
     * 通用人脸识别引擎对象池
     */
    private GenericObjectPool<FaceEngine> faceEngineGenericObjectPool;

    /**
     * 人脸比对引擎池
     */
    private GenericObjectPool<FaceEngine> faceEngineComparePool;

    @PostConstruct
    public void init(){
        GenericObjectPoolConfig detectPoolConfig = new GenericObjectPoolConfig();
        detectPoolConfig.setMaxIdle(5);
        detectPoolConfig.setMaxTotal(5);
        detectPoolConfig.setMinIdle(5);
        detectPoolConfig.setLifo(false);

        // 功能配置
        FunctionConfiguration detectFunctionCfg = new FunctionConfiguration();
        //开启人脸检测功能
        detectFunctionCfg.setSupportFaceDetect(true);
        //开启人脸识别功能
        detectFunctionCfg.setSupportFaceRecognition(true);
        //开启年龄检测功能
        detectFunctionCfg.setSupportAge(true);
        //开启性别检测功能
        detectFunctionCfg.setSupportGender(true);
        //开启活体检测功能
        detectFunctionCfg.setSupportLiveness(true);

        // 引擎配置
        EngineConfiguration detectCfg = new EngineConfiguration();
        //图片检测模式，如果是连续帧的视频流图片，那么改成VIDEO模式
        detectCfg.setDetectMode(DetectMode.ASF_DETECT_MODE_IMAGE);
        detectCfg.setFunctionConfiguration(detectFunctionCfg);
        //人脸旋转角度
        detectCfg.setDetectFaceOrientPriority(DetectOrient.ASF_OP_0_ONLY);
        //底层库算法对象池
        faceEngineGenericObjectPool = new GenericObjectPool(new FaceEngineFactory(arcsoftProperties.getLibpath(), arcsoftProperties.getAppid(), arcsoftProperties.getSdkkey(), detectCfg), detectPoolConfig);

        //初始化特征比较线程池
        GenericObjectPoolConfig comparePoolConfig = new GenericObjectPoolConfig();
        comparePoolConfig.setMaxIdle(5);
        comparePoolConfig.setMaxTotal(5);
        comparePoolConfig.setMinIdle(0);
        comparePoolConfig.setLifo(false);
        EngineConfiguration compareCfg = new EngineConfiguration();
        FunctionConfiguration compareFunctionCfg = new FunctionConfiguration();
        //开启人脸识别功能
        compareFunctionCfg.setSupportFaceRecognition(true);
        compareCfg.setFunctionConfiguration(compareFunctionCfg);
        //图片检测模式，如果是连续帧的视频流图片，那么改成VIDEO模式
        compareCfg.setDetectMode(DetectMode.ASF_DETECT_MODE_IMAGE);
        //人脸旋转角度
        compareCfg.setDetectFaceOrientPriority(DetectOrient.ASF_OP_0_ONLY);
        //底层库算法对象池
        faceEngineComparePool = new GenericObjectPool(new FaceEngineFactory(arcsoftProperties.getLibpath(), arcsoftProperties.getAppid(), arcsoftProperties.getSdkkey(),  compareCfg), comparePoolConfig);
        compareExecutorService = Executors.newFixedThreadPool(10);
    }

    @Override
    public List<FaceInfo> detectFaces(ImageInfo imageInfo) {
        FaceEngine faceEngine = null;
        try {
            // 从人脸引擎对象池中获取对象
            faceEngine = faceEngineGenericObjectPool.borrowObject();
            if (faceEngine == null) {
                // 抛出异常
                throw new CustomException(ErrorCodeEnum.BUSINESS_FAIL,"arcface engine get fail.","人脸识别引擎获取失败。");
            }
            //人脸检测得到人脸列表
            List<FaceInfo> faceInfoList = new ArrayList<FaceInfo>();
            //人脸检测
            int errorCode = faceEngine.detectFaces(imageInfo.getImageData(), imageInfo.getWidth(), imageInfo.getHeight(), imageInfo.getImageFormat(), faceInfoList);
            if (errorCode == 0) {
                return faceInfoList;
            } else {
                log.error("人脸检测失败，errorCode：" + errorCode);
            }

        } catch (Exception e) {
            log.error("", e);
        } finally {
            if (faceEngine != null) {
                //释放引擎对象
                faceEngineGenericObjectPool.returnObject(faceEngine);
            }
        }

        return null;
    }

    @Override
    public byte[] extractFaceFeature(ImageInfo imageInfo, FaceInfo faceInfo) {
        FaceEngine faceEngine = null;
        try {
            faceEngine = faceEngineGenericObjectPool.borrowObject();
            if (faceEngine == null) {
                throw new CustomException(ErrorCodeEnum.BUSINESS_FAIL,"arcface engine get fail.","人脸识别引擎获取失败。");
            }

            FaceFeature faceFeature = new FaceFeature();
            //提取人脸特征
            int errorCode = faceEngine.extractFaceFeature(imageInfo.getImageData(), imageInfo.getWidth(), imageInfo.getHeight(), imageInfo.getImageFormat(), faceInfo, faceFeature);
            if (errorCode == 0) {
                return faceFeature.getFeatureData();
            } else {
                log.error("特征提取失败，errorCode：" + errorCode);
            }

        } catch (Exception e) {
            log.error("", e);
        } finally {
            if (faceEngine != null) {
                //释放引擎对象
                faceEngineGenericObjectPool.returnObject(faceEngine);
            }
        }

        return null;
    }

    @Override
    public List<UserCompareInfo> faceRecognition(byte[] faceFeature, List<FacePerson> userInfoList, float passRate) {
        //识别到的人脸列表
        List<UserCompareInfo> resultUserInfoList = Lists.newLinkedList();

        FaceFeature targetFaceFeature = new FaceFeature();
        targetFaceFeature.setFeatureData(faceFeature);
        //分成1000一组，多线程处理
        List<List<FacePerson>> faceUserInfoPartList = Lists.partition(userInfoList, 1000);
        CompletionService<List<UserCompareInfo>> completionService = new ExecutorCompletionService(compareExecutorService);
        for (List<FacePerson> part : faceUserInfoPartList) {
            completionService.submit(new CompareFaceTask(part, targetFaceFeature, passRate));
        }
        for (int i = 0; i < faceUserInfoPartList.size(); i++) {
            List<UserCompareInfo> faceUserInfoList = null;
            try {
                faceUserInfoList = completionService.take().get();
            } catch (InterruptedException | ExecutionException e) {
            }
            if (CollectionUtil.isNotEmpty(userInfoList)) {
                resultUserInfoList.addAll(faceUserInfoList);
            }
        }
        //从大到小排序
        resultUserInfoList.sort((h1, h2) -> h2.getSimilar().compareTo(h1.getSimilar()));

        return resultUserInfoList;
    }

    private class CompareFaceTask implements Callable<List<UserCompareInfo>> {

        private List<FacePerson> facePeopleList;
        private FaceFeature targetFaceFeature;
        private float passRate;


        public CompareFaceTask(List<FacePerson> facePeopleList, FaceFeature targetFaceFeature, float passRate) {
            this.facePeopleList = facePeopleList;
            this.targetFaceFeature = targetFaceFeature;
            this.passRate = passRate;
        }

        @Override
        public List<UserCompareInfo> call() {
            FaceEngine faceEngine = null;
            List<UserCompareInfo> resultUserInfoList = Lists.newLinkedList();//识别到的人脸列表
            try {
                faceEngine = faceEngineComparePool.borrowObject();
                // 比对人脸
                for (FacePerson facePerson : facePeopleList) {
                    FaceFeature sourceFaceFeature = new FaceFeature();
                    sourceFaceFeature.setFeatureData(facePerson.getFeature());
                    FaceSimilar faceSimilar = new FaceSimilar();
                    faceEngine.compareFaceFeature(targetFaceFeature, sourceFaceFeature, faceSimilar);
                    if (faceSimilar.getScore() > passRate) {//相似值大于配置预期，加入到识别到人脸的列表
                        UserCompareInfo info = new UserCompareInfo();
                        info.setName(facePerson.getName());
                        info.setId(facePerson.getId());
                        info.setSimilar(faceSimilar.getScore());
                        resultUserInfoList.add(info);
                    }
                }
            } catch (Exception e) {
                log.error("", e);
            } finally {
                if (faceEngine != null) {
                    faceEngineComparePool.returnObject(faceEngine);
                }
            }
            return resultUserInfoList;
        }

    }


}
